package com.example.webcrawler.service.mpl;

import org.bouncycastle.asn1.x509.*;
import org.bouncycastle.asn1.x509.BasicConstraints;
import org.bouncycastle.asn1.x509.Extension;
import org.bouncycastle.asn1.x509.GeneralName;
import org.bouncycastle.asn1.x509.GeneralNames;
import org.bouncycastle.asn1.x509.KeyPurposeId;
import org.bouncycastle.asn1.x509.KeyUsage;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509ExtensionUtils;
import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.bouncycastle.asn1.x500.X500Name;
import org.springframework.stereotype.Service;

import javax.net.ssl.*;
import java.io.*;
import java.math.BigInteger;
import java.net.*;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.security.spec.RSAPrivateKeySpec;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@Service
public class HttpsProxyService {
    private static final Logger logger = LoggerFactory.getLogger(HttpsProxyService.class);

    // 存储已生成的站点证书
    private final Map<String, KeyPair> siteKeyPairs = new ConcurrentHashMap<>();
    private final Map<String, X509Certificate> siteCertificates = new ConcurrentHashMap<>();

    // CA证书和私钥
    private X509Certificate caCert;
    private PrivateKey caPrivateKey;

    // 代理服务器相关
    private ServerSocket proxyServerSocket;
    private ExecutorService executorService;
    private volatile boolean isRunning = false;
    private int proxyPort = 8081; // 默认代理端口

    public HttpsProxyService() {
        initCA();
    }

    /**
     * 启动HTTPS代理服务器（使用默认端口8081）
     */
    public void startProxy() throws IOException {
        startProxy(8081);
    }

    /**
     * 初始化CA证书
     */
    private void initCA() {
        try {
            // 检查是否存在已保存的CA证书
            File certFile = new File("webcrawler-ca.crt");
            File keyFile = new File("webcrawler-ca.key");

            if (certFile.exists() && keyFile.exists()) {
                loadCAFromFile(certFile, keyFile);
            } else {
                generateCACertificate();
                saveCAToFile(certFile, keyFile);
            }

            // 验证CA证书
            if (caCert != null) {
                caCert.checkValidity();
            }
        } catch (Exception e) {
            logger.error("初始化CA证书失败: {}", e.getMessage(), e);
            // 即使失败也生成一个新的
            try {
                generateCACertificate();
            } catch (Exception ex) {
                logger.error("生成CA证书失败: {}", ex.getMessage(), ex);
            }
        }
    }
    /**
     * 从文件加载CA证书
     */
    private void loadCAFromFile(File certFile, File keyFile) throws Exception {
        logger.info("从文件加载CA证书: {}", certFile.getAbsolutePath());
        // 简化实现，实际应该从文件读取证书和私钥
        generateCACertificate();
    }

    /**
     * 保存CA证书到文件
     */
    private void saveCAToFile(File certFile, File keyFile) {
        try (FileOutputStream certOut = new FileOutputStream(certFile);
             FileOutputStream keyOut = new FileOutputStream(keyFile)) {

            // 写入证书 (PEM格式)
            certOut.write("-----BEGIN CERTIFICATE-----\n".getBytes());
            certOut.write(Base64.getEncoder().encode(caCert.getEncoded()));
            certOut.write("\n-----END CERTIFICATE-----\n".getBytes());

            // 写入私钥 (PEM格式)
            KeyFactory keyFact = KeyFactory.getInstance("RSA");
            RSAPrivateKeySpec keySpec = keyFact.getKeySpec(caPrivateKey, RSAPrivateKeySpec.class);
            keyOut.write("-----BEGIN PRIVATE KEY-----\n".getBytes());
            keyOut.write(Base64.getEncoder().encode(keySpec.getModulus().toByteArray()));
            keyOut.write("\n-----END PRIVATE KEY-----\n".getBytes());

            logger.info("CA证书已保存到: {}", certFile.getAbsolutePath());
            logger.info("CA私钥已保存到: {}", keyFile.getAbsolutePath());
        } catch (Exception e) {
            logger.error("保存CA证书失败: {}", e.getMessage(), e);
        }
    }

    /**
     * 生成CA证书
     */
    private void generateCACertificate() throws Exception {
        // 生成密钥对
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();

        // 获取公钥和私钥
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();

        // 创建证书信息
        X500Name issuer = new X500Name("CN=WebCrawler CA, OU=WebCrawler Proxy, O=WebCrawler, L=Internet, ST=Internet, C=CN");
        BigInteger serial = BigInteger.valueOf(new SecureRandom().nextLong());
        Date notBefore = new Date(System.currentTimeMillis() - 24 * 60 * 60 * 1000); // 1天前
        Date notAfter = new Date(System.currentTimeMillis() + 10L * 365 * 24 * 60 * 60 * 1000); // 10年后
        X500Name subject = issuer;

        // 创建证书
        X509v3CertificateBuilder certificateBuilder = new JcaX509v3CertificateBuilder(
                issuer, serial, notBefore, notAfter, subject, publicKey);

        // 设置基本约束（CA证书）
        certificateBuilder.addExtension(Extension.basicConstraints, true, new BasicConstraints(true));

        // 设置密钥用法
        certificateBuilder.addExtension(Extension.keyUsage, true,
                new KeyUsage(KeyUsage.keyCertSign | KeyUsage.cRLSign | KeyUsage.digitalSignature));

        // 设置扩展密钥用法（客户端认证和服务器认证）
        certificateBuilder.addExtension(Extension.extendedKeyUsage, false,
                new org.bouncycastle.asn1.x509.ExtendedKeyUsage(new KeyPurposeId[]{
                        KeyPurposeId.id_kp_clientAuth,
                        KeyPurposeId.id_kp_serverAuth
                }));

        // 设置主体密钥标识符
        JcaX509ExtensionUtils extUtils = new JcaX509ExtensionUtils();
        certificateBuilder.addExtension(Extension.subjectKeyIdentifier, false,
                extUtils.createSubjectKeyIdentifier(publicKey));

        // 签名证书
        ContentSigner contentSigner = new JcaContentSignerBuilder("SHA256WithRSA")
                .build(privateKey);
        X509CertificateHolder certificateHolder = certificateBuilder.build(contentSigner);

        // 转换为X509Certificate
        JcaX509CertificateConverter certificateConverter = new JcaX509CertificateConverter();
        X509Certificate certificate = certificateConverter.getCertificate(certificateHolder);

        // 验证证书
        certificate.checkValidity();
        certificate.verify(publicKey);

        // 保存证书和私钥
        this.caCert = certificate;
        this.caPrivateKey = privateKey;

        logger.info("CA证书生成成功");
    }

    /**
     * 为特定主机生成证书
     */
    private X509Certificate generateCertificateForHost(String host) throws Exception {
        logger.debug("为站点 {} 生成证书", host);

        // 获取或创建密钥对
        KeyPair siteKeyPair = siteKeyPairs.get(host);
        if (siteKeyPair == null) {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
            keyGen.initialize(2048);
            siteKeyPair = keyGen.generateKeyPair();
            siteKeyPairs.put(host, siteKeyPair);
        }

        // 创建证书信息
        X500Name issuer = new X500Name(caCert.getSubjectDN().getName());
        BigInteger serial = BigInteger.valueOf(new SecureRandom().nextLong());
        Date notBefore = new Date(System.currentTimeMillis() - 60 * 60 * 1000); // 1小时前
        Date notAfter = new Date(System.currentTimeMillis() + 365L * 24 * 60 * 60 * 1000); // 1年
        X500Name subject = new X500Name("CN=" + host + ", OU=WebCrawler Proxy, O=WebCrawler, L=Internet, ST=Internet, C=CN");

        // 创建证书
        X509v3CertificateBuilder certificateBuilder = new JcaX509v3CertificateBuilder(
                issuer, serial, notBefore, notAfter, subject, siteKeyPair.getPublic());

        // 设置基本约束（非CA证书）
        certificateBuilder.addExtension(Extension.basicConstraints, true, new BasicConstraints(false));

        // 设置密钥用法
        certificateBuilder.addExtension(Extension.keyUsage, true,
                new KeyUsage(KeyUsage.digitalSignature | KeyUsage.keyEncipherment | KeyUsage.keyAgreement));

        // 设置扩展密钥用法（服务器认证）
        certificateBuilder.addExtension(Extension.extendedKeyUsage, false,
                new org.bouncycastle.asn1.x509.ExtendedKeyUsage(
                        new org.bouncycastle.asn1.x509.KeyPurposeId[]{
                                org.bouncycastle.asn1.x509.KeyPurposeId.id_kp_serverAuth
                        }));

        // 设置主体可选名称（SAN）
        GeneralName generalName = new GeneralName(GeneralName.dNSName, host);
        GeneralNames generalNames = new GeneralNames(generalName);
        certificateBuilder.addExtension(Extension.subjectAlternativeName, false, generalNames);

        // 设置颁发者密钥标识符
        JcaX509ExtensionUtils extUtils = new JcaX509ExtensionUtils();
        certificateBuilder.addExtension(Extension.authorityKeyIdentifier, false,
                extUtils.createAuthorityKeyIdentifier(caCert.getPublicKey()));

        // 设置主体密钥标识符
        certificateBuilder.addExtension(Extension.subjectKeyIdentifier, false,
                extUtils.createSubjectKeyIdentifier(siteKeyPair.getPublic()));

        // 签名证书
        ContentSigner contentSigner = new JcaContentSignerBuilder("SHA256WithRSA")
                .build(caPrivateKey);
        X509Certificate certificate = new JcaX509CertificateConverter()
                .getCertificate(certificateBuilder.build(contentSigner));

        // 验证证书
        certificate.checkValidity();
        certificate.verify(caCert.getPublicKey());

        // 缓存证书
        siteCertificates.put(host, certificate);

        logger.debug("站点 {} 的证书生成成功", host);
        return certificate;
    }

    /**
     * 启动HTTPS代理服务器
     */
    public void startProxy(int port) throws IOException {
        if (isRunning) {
            logger.warn("代理服务器已在运行中");
            return;
        }

        this.proxyPort = port;
        proxyServerSocket = new ServerSocket(proxyPort);
        executorService = Executors.newCachedThreadPool();
        isRunning = true;

        logger.info("HTTPS代理服务器启动，监听端口: {}", port);

        // 启动代理服务器线程
        executorService.submit(() -> {
            while (isRunning && !proxyServerSocket.isClosed()) {
                try {
                    Socket clientSocket = proxyServerSocket.accept();
                    executorService.submit(() -> handleClientRequest(clientSocket));
                } catch (IOException e) {
                    if (isRunning) {
                        logger.error("接受客户端连接时出错: {}", e.getMessage(), e);
                    }
                }
            }
        });
    }

    /**
     * 停止HTTPS代理服务器
     */
    public void stopProxy() {
        if (!isRunning) {
            logger.warn("代理服务器未在运行");
            return;
        }

        isRunning = false;

        try {
            if (proxyServerSocket != null && !proxyServerSocket.isClosed()) {
                proxyServerSocket.close();
            }
        } catch (IOException e) {
            logger.error("关闭代理服务器套接字时出错: {}", e.getMessage(), e);
        }

        if (executorService != null) {
            executorService.shutdownNow();
        }

        logger.info("HTTPS代理服务器已停止");
    }

    /**
     * 处理客户端请求
     */
    private void handleClientRequest(Socket clientSocket) {
        try {
            InputStream clientIn = clientSocket.getInputStream();
            OutputStream clientOut = clientSocket.getOutputStream();

            // 读取客户端请求
            BufferedReader reader = new BufferedReader(new InputStreamReader(clientIn));
            String requestLine = reader.readLine();

            if (requestLine == null) {
                clientSocket.close();
                return;
            }

            logger.debug("接收到请求: {}", requestLine);

            // 解析请求行
            String[] parts = requestLine.split(" ");
            if (parts.length < 3) {
                sendError(clientOut, "400 Bad Request");
                clientSocket.close();
                return;
            }

            String method = parts[0];
            String target = parts[1];
            String version = parts[2];

            // 处理CONNECT方法（HTTPS流量）
            if ("CONNECT".equals(method)) {
                handleHttpsConnect(clientSocket, target);
            } else {
                // 处理普通HTTP请求
                handleHttpRegular(clientSocket, requestLine, reader);
            }
        } catch (Exception e) {
            logger.error("处理客户端请求时出错: {}", e.getMessage(), e);
            try {
                clientSocket.close();
            } catch (IOException ioException) {
                logger.error("关闭客户端套接字时出错: {}", ioException.getMessage(), ioException);
            }
        }
    }

    /**
     * 处理HTTPS CONNECT请求
     */
    private void handleHttpsConnect(Socket clientSocket, String target) {
        try {
            String[] hostPort = target.split(":");
            String host = hostPort[0];
            int port = hostPort.length > 1 ? Integer.parseInt(hostPort[1]) : 443;

            logger.debug("处理HTTPS连接到 {}:{}", host, port);

            // 连接到目标服务器
            Socket targetSocket = new Socket(host, port);

            // 发送连接成功的响应给客户端
            OutputStream clientOut = clientSocket.getOutputStream();
            clientOut.write("HTTP/1.1 200 Connection Established\r\n\r\n".getBytes());
            clientOut.flush();

            // 执行SSL握手并拦截流量
            interceptHttpsTraffic(clientSocket, targetSocket, host);

        } catch (Exception e) {
            logger.error("处理HTTPS CONNECT请求时出错: {}", e.getMessage(), e);
            try {
                OutputStream clientOut = clientSocket.getOutputStream();
                sendError(clientOut, "500 Internal Server Error");
                clientSocket.close();
            } catch (IOException ioException) {
                logger.error("发送错误响应时出错: {}", ioException.getMessage(), ioException);
            }
        }
    }

    /**
     * 拦截HTTPS流量
     */
    private void interceptHttpsTraffic(Socket clientSocket, Socket targetSocket, String host) {
        try {
            // 创建SSL上下文
            SSLContext sslContext = SSLContext.getInstance("TLS");

            // 为站点生成证书
            X509Certificate siteCert = siteCertificates.get(host);
            if (siteCert == null) {
                siteCert = generateCertificateForHost(host);
            }

            if (siteCert == null) {
                throw new Exception("无法为站点 " + host + " 生成证书");
            }

            // 获取站点私钥
            KeyPair siteKeyPair = siteKeyPairs.get(host);
            if (siteKeyPair == null) {
                throw new Exception("未找到站点 " + host + " 的私钥");
            }

            // 初始化SSL上下文
            KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            KeyStore keyStore = KeyStore.getInstance("JKS");
            keyStore.load(null, null);

            // 正确设置证书链，CA证书应该在站点证书之后
            Certificate[] certChain = new Certificate[]{siteCert, caCert};
            char[] password = "changeit".toCharArray();
            keyStore.setKeyEntry("alias", siteKeyPair.getPrivate(), password, certChain);
            kmf.init(keyStore, password);

            // 创建自定义的TrustManager，接受所有证书（仅用于代理目的）
            TrustManager[] trustAllCerts = new TrustManager[] {
                    new X509TrustManager() {
                        @Override
                        public X509Certificate[] getAcceptedIssuers() {
                            return new X509Certificate[0];
                        }
                        @Override
                        public void checkClientTrusted(X509Certificate[] certs, String authType) {
                        }
                        @Override
                        public void checkServerTrusted(X509Certificate[] certs, String authType) {
                        }
                    }
            };

            sslContext.init(kmf.getKeyManagers(), trustAllCerts, new SecureRandom());

            // 执行SSL握手
            SSLSocketFactory sslSocketFactory = sslContext.getSocketFactory();
            SSLSocket sslClientSocket = (SSLSocket) sslSocketFactory.createSocket(
                    clientSocket, clientSocket.getInetAddress().getHostAddress(), clientSocket.getPort(), true);

            // 设置支持的加密套件和协议
            sslClientSocket.setEnabledCipherSuites(sslSocketFactory.getSupportedCipherSuites());
            sslClientSocket.setEnabledProtocols(new String[]{"TLSv1.2", "TLSv1.3"});
            sslClientSocket.setUseClientMode(false);
            sslClientSocket.setNeedClientAuth(false);

            // 增加握手超时时间
            sslClientSocket.setSoTimeout(30000);
            sslClientSocket.startHandshake();

            // 建立到目标服务器的SSL连接
            SSLSocket sslTargetSocket = (SSLSocket) sslSocketFactory.createSocket(
                    targetSocket, targetSocket.getInetAddress().getHostAddress(), targetSocket.getPort(), true);
            sslTargetSocket.setEnabledCipherSuites(sslSocketFactory.getSupportedCipherSuites());
            sslTargetSocket.setEnabledProtocols(new String[]{"TLSv1.2", "TLSv1.3"});
            sslTargetSocket.setUseClientMode(true);

            // 增加握手超时时间
            sslTargetSocket.setSoTimeout(30000);
            sslTargetSocket.startHandshake();

            // 转发数据（拦截在此处实现）
            forwardData(sslClientSocket, sslTargetSocket, host);

        } catch (Exception e) {
            logger.error("拦截HTTPS流量时出错: {}", e.getMessage(), e);
            try {
                clientSocket.close();
                targetSocket.close();
            } catch (IOException ioException) {
                logger.error("关闭套接字时出错: {}", ioException.getMessage(), ioException);
            }
        }
    }

    /**
     * 转发数据并在过程中拦截
     */
    private void forwardData(SSLSocket clientSocket, SSLSocket targetSocket, String host) {
        try {
            // 创建两个线程来转发两个方向的数据
            Thread clientToServer = new Thread(() -> {
                try {
                    InputStream in = clientSocket.getInputStream();
                    OutputStream out = targetSocket.getOutputStream();
                    byte[] buffer = new byte[8192];
                    int len;
                    while ((len = in.read(buffer)) != -1) {
                        // 在此处可以拦截和分析请求数据
                        String requestData = new String(buffer, 0, len);
                        logger.debug("客户端 -> 服务器 数据: {}", requestData);

                        // 可以在这里对请求数据进行分析和记录
                        // 例如：解析HTTP请求头、记录请求内容等

                        out.write(buffer, 0, len);
                        out.flush();
                    }
                } catch (IOException e) {
                    logger.debug("客户端到服务器数据转发结束: {}", e.getMessage());
                }
            });

            Thread serverToClient = new Thread(() -> {
                try {
                    InputStream in = targetSocket.getInputStream();
                    OutputStream out = clientSocket.getOutputStream();
                    byte[] buffer = new byte[8192];
                    int len;
                    while ((len = in.read(buffer)) != -1) {
                        // 在此处可以拦截和分析响应数据
                        String responseData = new String(buffer, 0, len);
                        logger.debug("服务器 -> 客户端 数据: {}", responseData);

                        // 可以在这里对响应数据进行分析和记录
                        // 例如：解析HTTP响应头、记录响应内容等

                        out.write(buffer, 0, len);
                        out.flush();
                    }
                } catch (IOException e) {
                    logger.debug("服务器到客户端数据转发结束: {}", e.getMessage());
                }
            });

            clientToServer.start();
            serverToClient.start();

            clientToServer.join();
            serverToClient.join();

        } catch (Exception e) {
            logger.error("转发数据时出错: {}", e.getMessage(), e);
        } finally {
            try {
                clientSocket.close();
                targetSocket.close();
            } catch (IOException e) {
                logger.error("关闭SSL套接字时出错: {}", e.getMessage(), e);
            }
        }
    }

    /**
     * 处理普通HTTP请求
     */
    private void handleHttpRegular(Socket clientSocket, String requestLine, BufferedReader reader) throws IOException {
        // 对于普通HTTP请求，直接转发
        OutputStream clientOut = clientSocket.getOutputStream();
        clientOut.write("HTTP/1.1 200 OK\r\n".getBytes());
        clientOut.write("Content-Length: 2\r\n".getBytes());
        clientOut.write("\r\n".getBytes());
        clientOut.write("OK".getBytes());
        clientOut.flush();
        clientSocket.close();
    }

    /**
     * 发送错误响应
     */
    private void sendError(OutputStream out, String error) throws IOException {
        String response = "HTTP/1.1 " + error + "\r\n" +
                "Content-Length: 0\r\n" +
                "\r\n";
        out.write(response.getBytes());
        out.flush();
    }

    /**
     * 获取CA证书，用于安装到系统/浏览器信任存储区
     */
    public X509Certificate getCACertificate() {
        return caCert;
    }

    /**
     * 检查代理是否正在运行
     */
    public boolean isProxyRunning() {
        return isRunning;
    }

    /**
     * 获取代理端口
     */
    public int getProxyPort() {
        return proxyPort;
    }

    /**
     * 获取CA证书Base64编码字符串，用于安装到系统/浏览器信任存储区
     */
    public String getCACertificateBase64() {
        try {
            if (caCert != null) {
                return Base64.getEncoder().encodeToString(caCert.getEncoded());
            }
        } catch (Exception e) {
            logger.error("编码CA证书时出错: {}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 获取CA证书PEM格式字符串，用于安装到系统/浏览器信任存储区
     */
    public String getCACertificatePEM() {
        try {
            if (caCert != null) {
                StringBuilder pem = new StringBuilder();
                pem.append("-----BEGIN CERTIFICATE-----\n");
                pem.append(Base64.getEncoder().encodeToString(caCert.getEncoded()));
                pem.append("\n-----END CERTIFICATE-----\n");
                return pem.toString();
            }
        } catch (Exception e) {
            logger.error("生成CA证书PEM时出错: {}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 将CA证书保存为PEM文件
     */
    public void saveCACertificateAsPEM(String filename) {
        try {
            String pem = getCACertificatePEM();
            if (pem != null) {
                try (FileWriter writer = new FileWriter(filename)) {
                    writer.write(pem);
                }
                logger.info("CA证书已保存为PEM格式: {}", filename);
            }
        } catch (Exception e) {
            logger.error("保存CA证书为PEM格式时出错: {}", e.getMessage(), e);
        }
    }
}
